package main

import (
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"time"

	"github.com/gorilla/websocket"
	"golang.org/x/crypto/ssh"
)

var upgrader = websocket.Upgrader{
	CheckOrigin: func(r *http.Request) bool {
		return true // 允许所有跨域请求
	},
}

type SSHParams struct {
	Host     string `json:"host"`
	Username string `json:"username"`
	Password string `json:"password"`
	Script   string `json:"script"`
}

type Message struct {
	Type string `json:"type"`
	Data string `json:"data"`
}

func handleWebSocket(w http.ResponseWriter, r *http.Request) {
	// 升级HTTP连接为WebSocket
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Printf("升级WebSocket失败: %v", err)
		return
	}
	defer conn.Close()

	// 读取参数
	_, msg, err := conn.ReadMessage()
	if err != nil {
		sendError(conn, "读取参数失败: "+err.Error())
		return
	}

	var params SSHParams
	if err := json.Unmarshal(msg, &params); err != nil {
		sendError(conn, "参数解析失败: "+err.Error())
		return
	}

	// 发送连接消息
	sendMessage(conn, "正在连接到服务器 "+params.Host+"...")

	// 创建SSH配置
	config := &ssh.ClientConfig{
		User: params.Username,
		Auth: []ssh.AuthMethod{
			ssh.Password(params.Password),
		},
		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
		Timeout:         30 * time.Second,
	}

	// 连接SSH
	client, err := ssh.Dial("tcp", fmt.Sprintf("%s:22", params.Host), config)
	if err != nil {
		sendError(conn, "SSH连接失败: "+err.Error())
		return
	}
	defer client.Close()

	sendMessage(conn, "SSH连接已建立")

	// 创建会话
	session, err := client.NewSession()
	if err != nil {
		sendError(conn, "创建会话失败: "+err.Error())
		return
	}
	defer session.Close()

	// 请求伪终端
	if err := session.RequestPty("xterm", 40, 80, ssh.TerminalModes{}); err != nil {
		sendError(conn, "请求PTY失败: "+err.Error())
		return
	}

	// 创建管道接收输出
	stdout, err := session.StdoutPipe()
	if err != nil {
		sendError(conn, "创建输出管道失败: "+err.Error())
		return
	}

	stderr, err := session.StderrPipe()
	if err != nil {
		sendError(conn, "创建错误输出管道失败: "+err.Error())
		return
	}

	// 启动命令
	command := fmt.Sprintf("python3 -u %s 2>&1", params.Script)
	if err := session.Start(command); err != nil {
		sendError(conn, "启动命令失败: "+err.Error())
		return
	}

	// 读取输出
	go func() {
		buffer := make([]byte, 1024)
		for {
			n, err := stdout.Read(buffer)
			if n > 0 {
				sendMessage(conn, string(buffer[:n]))
			}
			if err != nil {
				break
			}
		}
	}()

	// 读取错误输出
	go func() {
		buffer := make([]byte, 1024)
		for {
			n, err := stderr.Read(buffer)
			if n > 0 {
				sendError(conn, string(buffer[:n]))
			}
			if err != nil {
				break
			}
		}
	}()

	// 等待命令完成
	if err := session.Wait(); err != nil {
		sendError(conn, "命令执行错误: "+err.Error())
		return
	}

	sendMessage(conn, "脚本执行完毕")
}

func sendMessage(conn *websocket.Conn, data string) {
	msg := Message{
		Type: "message",
		Data: data,
	}
	if err := conn.WriteJSON(msg); err != nil {
		log.Printf("发送消息失败: %v", err)
	}
}

func sendError(conn *websocket.Conn, data string) {
	msg := Message{
		Type: "error",
		Data: data,
	}
	if err := conn.WriteJSON(msg); err != nil {
		log.Printf("发送错误消息失败: %v", err)
	}
}

func main() {
	http.HandleFunc("/ws", handleWebSocket)
	log.Printf("服务器启动在 :2345")
	log.Fatal(http.ListenAndServe(":2345", nil))
}
